##########################################################################################

library(data.table)
library(optparse)
library(ArchR)
library(dplyr)
library(chromVAR)
library(BiocParallel)

##########################################################################################
option_list <- list(
    make_option(c("--comine_data_file"), type = "character"),
    make_option(c("--scriptPath"), type = "character"),
    make_option(c("--cpu"), type = "character"),
    make_option(c("--use_tf_file"), type = "character"),
    make_option(c("--cluster"), type = "character"),
    make_option(c("--out_path"), type = "character") 
)

if(1!=1){
    
    ## 整合atac和rna的文件
    comine_data_file <- "~/20231121_singleMuti/results/subcell/cluster5/cluster5.combineRNA.motif_peak2gene.Rdata"

    ## cluster
    cluster <- "cluster5"

    ## tf_file
    use_tf_file <- "~/20231121_singleMuti/results/celltype_plot/motif_motif/use_tf.list"

    cpu <- 12

    ## 既往研究整理的代码
    scriptPath <- "~/20231121_singleMuti/scripts/scScalpChromatin"

    ## 输出
    out_path <- "~/20231121_singleMuti/results/celltype_plot/motif_motif"

}

###########################################################################################
parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

comine_data_file <- opt$comine_data_file
cluster <- opt$cluster
out_path <- opt$out_path
use_tf_file <- opt$use_tf_file
scriptPath <- opt$scriptPath
cpu <- as.numeric(opt$cpu)

dir.create(out_path , recursive = T)

###########################################################################################

#addArchRThreads(threads = 10)
#addArchRGenome("hg38") 
## 设置chromvar的线程
#cpu <- 10
register(MulticoreParam(cpu))

## 已发表文献写好的脚本
source(paste0(scriptPath, "/misc_helpers.R"))
source(paste0(scriptPath, "/matrix_helpers.R"))

###########################################################################################
#### 导入数据
a <- load(comine_data_file)
## atac_proj
if( cluster %in% c("all" , "somatic" , "germ") ){
    atac_proj <- testis_combined_peak_combineRNA
    #atac_proj <- addPeak2GeneLinks(ArchRProj = atac_proj , useMatrix = "GeneExpressionMatrix")
}

use_tf <- data.frame(fread(use_tf_file , header = F))$V1

###########################################################################################
#### 计算motif和motif的相关系数，自己写的代码
## motfi
motif_se <- getMatrixFromProject(atac_proj, useMatrix="MotifMatrix")
## The deviations are the bias corrected deviations in accessibility. For each motif or annotation (rows), 
## there is a value for each cell or sample (columns) representing how different the accessibility for peaks with that motif or annotation 
## is from the expectation based on all cells being equal, corrected for biases.
motif_mat <- assays(motif_se)$deviation
## 转化为数值矩阵
motif_mat_num <- apply(motif_mat , 1 , as.numeric)
rownames(motif_mat_num) <- colnames(motif_mat)
motif_mat_num <- t(motif_mat_num)

## 计算相关系数，默认的person相关系数
cor_result <- cor2Matrices(motif_mat_num , motif_mat_num)
## 输出
out_file <- paste0( out_path , "/" , cluster , ".motif_cor.tsv" )
write.table( cor_result , out_file , row.names = F , sep = "\t" )

###########################################################################################
#### chromvar自带的计算相关系数以及协同作用
## Get Peak Matrix
proj <- atac_proj
sePeak <- getMatrixFromProject(proj, "PeakMatrix")
names(assays(sePeak)) <- "counts" #Change for chromVAR
fragments_per_peak <- getFragmentsPerPeak(sePeak)
## 去除peak，没有细胞有counts
use_peak <- which(fragments_per_peak > 0)

## Set Bias
rowData(sePeak)$bias <- getPeakSet(proj)$GC
## Filtering peaks
sePeak <- filterPeaks(sePeak, non_overlapping = TRUE)

## Get Annotation Matrix
matches <- getMatches(proj)
matches <- matches[use_peak,]

## 提取感兴趣的TF的motif
use_tf_index <- which(sapply(strsplit(colnames(matches) , "_") , "[" , 1) %in% use_tf)
matches <- matches[,use_tf_index]

## Sanity Check
all(paste0(rowRanges(sePeak)) == paste0(rowRanges(matches)))

## Synergy
dat_synergy <- getAnnotationSynergy(sePeak, matches)
out_file <- paste0( out_path , "/" , cluster , ".motif_synergy-chromvar.tsv" )
write.table( dat_synergy , out_file , row.names = T , sep = "\t" )

## Correlatoin
dat_cor <- getAnnotationCorrelation(sePeak, matches)
out_file <- paste0( out_path , "/" , cluster , ".motif_cor-chromvar.tsv" )
write.table( dat_cor , out_file , row.names = T , sep = "\t" )
